#!/bin/env python

import sys
import subprocess
import argparse
import os
import re
import shutil

parser = argparse.ArgumentParser(description= """
DESCRIPTION
    Execute a series of steps starting from fastq files untrimmed and in sanger encoding.        
    The shell script goes in the same dir as the input fastq.
    MEMO: Currently pair-end reads must be in the same directory, only the first filename
        in the pair is used to generate the successive filenames.
    
EXAMPLE
    ## See what the pipeline consists of.
    ## It doesn't matter if the input file doesn't exists as --noexec will only
    ## create the script file:
    bismark_pipeline.py -f FASTQ -g GENOME -o OUTDIR --noexec
    cat FASTQ.sh
TODO:

""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('--fastq', '-f',
                   required= True,
                   nargs= '+',
                   help='''Fastq file(s) to process can be gzipped. It can be one file
or two files. If two files are passed they are treated as a pair.
                   ''')

parser.add_argument('--genome', '-g',
                   required= True,
                   help='''Reference genome directory for bismark. Use full path
                   ''')

parser.add_argument('--outdir', '-o',
                   required= False,
                   default= '.',
                   help='''Output directory. Default current wdir. In this dir a
directory called 'bismark-'<basename of fastq> will be created if it doesn't exists.
This subdir is necessary otherwise bismark_methylation_extractor will fail with
concurrent runs.
                   ''')

parser.add_argument('--sh',
                   action= 'store_true',
                   help='''Execute the commands on the front end.
With this option bsub is not used (which is the default).
                   ''')

parser.add_argument('--scriptname',
                   required= False,
                   default= None,
                   help='''The pipeline is written to this filename. Default is
<fastq>.sh
This file is then passed to bsub or executed on the current shell.
                   ''')

parser.add_argument('--mem',
                   type= int,
                   required= False,
                   default= 2048,
                   help='''Memory requirement for bsub -R "rusage[mem=xxxx]". Default is 2048.
                   ''')

parser.add_argument('--bsubOpt',
                   required= False,
                   default= None,
                   help='''Options to pass to bsub. Ignored if --sh flag is used.
NB: If this opiton is set the defualt (-J -oo --R) are not set. 
                   ''')

parser.add_argument('--CX',
                   action= 'store_true',
                   help='''Pass this option to bismark_methylation_extractor. It will
produce a pileup for each C in the genome
                   ''')

parser.add_argument('--noexec',
                   action= 'store_true',
                   help='''Produce the script files but don't execute them. The line that
is suppressed is printed to stdout.
                   ''')

parser.add_argument('--trim',
                   action= 'store_true',
                   help='''CURRENTLY DO NOT USE. Execute trim_galore on the input fastq
                   ''')

parser.add_argument('--rm_fastq',
                   action= 'store_true',
                   help='''USE WITH CARE: Removes the input fastq file once the pipeline
is completed. Useful to remove fastqfiles generated by split command.
                   ''')

parser.add_argument('--rm_trimmed_fq',
                   action= 'store_true',
                   help='''USE WITH CARE: Removes the trimmed fastq file once the pipeline
is completed. Useful to remove fastqfiles generated by split command.
                   ''')

args= parser.parse_args()

# -----------------------------------------------------------------------------

if args.trim:
    sys.exit('trimming not supported at the moment')

if len(args.fastq) == 1:
    paired= False
elif len(args.fastq) == 2:
    paired= True
else:
    sys.exit('Too many files passed to --fastq: %s' %(args.fastq))

if args.scriptname is None:
    scriptname= args.fastq[0] + '.sh'
else:
    scriptname= args.scriptname

if args.bsubOpt is None:
    bsubOpt= '-J bismark_pipeline-' + args.fastq[0] + ' -oo ' + args.fastq[0] + '.bsub.log' + ' -R "rusage[mem=%(mem)s]"' %{'mem': args.mem}
else:
    bsubOpt= args.bsubOpt
    
if args.fastq[0].endswith('.gz'):
    gz= '.gz'
else:
    gz= ''
# -----------------------------------------------------------------------------

## 
cwd= os.getcwd() 
fqInput= os.path.split(args.fastq[0])[1]
fqDir= os.path.split(args.fastq[0])[0]
if fqDir == '':
    fqDir= '.'
if fqInput.endswith('.fq.gz') or fqInput.endswith('.fq'):
    fqExt= '.fq'
elif fqInput.endswith('.fastq.gz') or fqInput.endswith('.fastq'):
    fqExt= '.fastq'
else:
    sys.exit('Invalid extension for %s: Must be .fq or .fastq with or without .gz' %args.fastq)
    
fqBname= re.sub('\.gz$', '', fqInput )
fqBname= re.sub('\.fastq$|\.fq$', '', fqBname)

## Full path to genome
genome= os.path.abspath(args.genome)

## Suffixes:
if args.trim:
    trim_suffix= '_trimmed.fq' + gz
else:
    trim_suffix= fqExt + gz
    
bismark_suffix= '_bt2_bismark.sam'
methylation_suffix= '_bt2_bismark.txt'

## Prepare output dir:
outdir= os.path.join(args.outdir, 'bismark-' + fqBname)

if args.CX:
    CX= '--CX'
else:
    CX= ''

if not os.path.exists(outdir):
    os.makedirs(outdir)
# -----------------------[ Compile commands ]----------------------------------
mkdir_cmd= 'mkdir -p %s' %(outdir)

if args.trim:
    trim_cmd= 'trim_galore -q 20 %(fqDir)s/%(fqInput)s; mv %(fqDir)s/%(fqInput)s_trimming_report.txt %(outdir)s' %{'fqDir': fqDir, 'fqInput': fqInput, 'outdir': outdir}
else:
    trim_cmd= ''

if paired:
    bismark_cmd= 'bismark --bowtie2 -o %(outdir)s %(genome)s -1 %(fq1)s -2 %(fq2)s' %{'outdir': outdir, 'genome': genome, 'fq1': args.fastq[0], 'fq2': args.fastq[1]}
else:
    bismark_cmd= 'bismark --bowtie2 -o %(outdir)s %(genome)s %(fqDir)s/%(fqBname)s%(suffix)s' %{'outdir': outdir, 'genome': genome, 'fqDir': fqDir, 'fqBname': fqBname, 'suffix': trim_suffix}

if args.rm_trimmed_fq:
    trim_rm= 'rm %(fqDir)s/%(fqBname)s%(suffix)s' %{'fqDir': fqDir, 'fqBname': fqBname, 'suffix': trim_suffix}
else:
    trim_rm= ''

methyl_cmd= 'cd %(outdir)s; bismark_methylation_extractor -o . --single-end --comprehensive --merge_non_CpG --report --cytosine_report %(CX)s --genome_folder %(genome)s ./%(fqBname)s%(suffix)s; cd %(cwd)s' \
            %{'outdir': outdir, 'genome': genome, 'fqBname': fqBname, 'suffix': trim_suffix + bismark_suffix, 'cwd': cwd, 'CX': CX}

methyl_bg_rm= 'rm %(outdir)s/%(fqBname)s%(suffix)s.bedGraph' %{'outdir': outdir, 'fqBname': fqBname, 'suffix': trim_suffix + re.sub('\.sam', '', bismark_suffix)}
methyl_noncpg_rm= 'rm %(outdir)s/Non_CpG_context_%(fqBname)s%(trim_suffix)s%(methylation_suffix)s' %{'outdir': outdir, 'fqBname': fqBname, 'trim_suffix':trim_suffix, 'methylation_suffix':methylation_suffix}
methyl_cpg_rm= 'rm %(outdir)s/CpG_context_%(fqBname)s%(trim_suffix)s%(methylation_suffix)s' %{'outdir': outdir, 'fqBname': fqBname, 'trim_suffix':trim_suffix, 'methylation_suffix':methylation_suffix}

sam2bam_cmd= '''samtools view -S -b %(outdir)s/%(fqBname)s%(suffix)s > %(outdir)s/%(fqBname)s.unsorted.bam &&
rm %(outdir)s/%(fqBname)s%(suffix)s &&
samtools sort %(outdir)s/%(fqBname)s.unsorted.bam %(outdir)s/%(fqBname)s &&
rm %(outdir)s/%(fqBname)s.unsorted.bam &&
samtools index %(outdir)s/%(fqBname)s.bam''' %{'outdir': outdir, 'fqBname': fqBname, 'suffix': trim_suffix + bismark_suffix}

if args.rm_fastq:
    rm_fastq_cmd= 'rm %s' %(args.fastq)
else:
    rm_fastq_cmd= 'echo'

# ------------------------------------------------------------------------------   
script= '\n\n'.join(['#!/bin/sh', 'set -e', 'echo "cd %s"' %(cwd),
                     mkdir_cmd,
                     trim_cmd,
                     bismark_cmd,
                     trim_rm,
                     methyl_cmd,
                     methyl_bg_rm,
                     methyl_noncpg_rm,
                     methyl_cpg_rm,
                     sam2bam_cmd,
                     rm_fastq_cmd,
                     'exit'])

fout= open(scriptname, 'w')
fout.write(script)
fout.close()

if args.sh:
   cmd= 'sh ' + scriptname
else:
    cmd= 'bsub ' + bsubOpt + ' < ' + scriptname

if args.noexec:
    print(cmd)
else:
    p= subprocess.Popen(cmd, shell= True)
    p.wait()
    print(cmd)
    cmd= 'mv %s %s/' %(scriptname, outdir)
    p= subprocess.Popen(cmd, shell= True)
    p.wait()
    print(cmd)
sys.exit()
